# %%
from torchvision import transforms as T
import matplotlib.pyplot as plt
from torchvision.datasets import ImageFolder

DATA_PATH = r'C:\files\git_repository\pytorch-learning\datasets\hymenoptera_data\train'
dataset = ImageFolder(DATA_PATH)


# %%
print(dataset.class_to_idx)
# 就是把文件夹名字转化成index
# {'train': 0, 'val': 1}
# %%
# 所有图片的路径和对应的label
print(dataset.imgs)
# %%
print(dataset.classes)
# %%
# 没有任何的transform，所以返回的还是PIL Image对象
print(dataset[0][1])  # 第一维是第几张图，第二维为1返回label
# print(dataset[0][0]) # 为0返回图片数据
plt.imshow(dataset[0][0])
plt.show()


# %%
normalize = T.Normalize(mean=[0.4, 0.4, 0.4], std=[0.2, 0.2, 0.2])
transform = T.Compose([
    T.RandomResizedCrop(224),
    T.RandomHorizontalFlip(),
    T.ToTensor(),
    normalize,
])
dataset = ImageFolder('data1/dogcat_2/', transform=transform)

# 深度学习中图片数据一般保存成CxHxW，即通道数x图片高x图片宽
# print(dataset[0][0].size())

to_img = T.ToPILImage()
# 0.2和0.4是标准差和均值的近似
a = to_img(dataset[0][0]*0.2+0.4)
plt.imshow(a)
plt.axis('off')
plt.show()
